package com.yuxl.admin.manager.auth.core.dataLoader;

import com.yuxl.admin.manager.auth.core.entity.AuthRole;
import com.yuxl.admin.manager.auth.core.entity.AuthUserRole;
import com.yuxl.admin.manager.auth.core.repository.AuthUserRoleRepository;
import com.yuxl.admin.manager.auth.core.repository.AuthRoleRepository;
import org.dataloader.BatchLoaderEnvironment;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Mono;

import java.util.List;
import java.util.Map;
import java.util.function.BiFunction;
import java.util.stream.Collectors;

@Component
public class UserRoleDataLoader implements BiFunction<List<Long>, BatchLoaderEnvironment, Mono<List<List<AuthRole>>>> {

    private final AuthUserRoleRepository userRoleRepository;
    private final AuthRoleRepository roleRepository;

    public UserRoleDataLoader(AuthUserRoleRepository userRoleRepository, AuthRoleRepository roleRepository) {
        this.userRoleRepository = userRoleRepository;
        this.roleRepository = roleRepository;
    }

    @Override
    public Mono<List<List<AuthRole>>> apply(List<Long> userIds, BatchLoaderEnvironment environment) {
        return userRoleRepository.findByUserIdIn(userIds)
                .collectList()
                .flatMap(userRoles -> {
                    List<Long> roleIds = userRoles.stream()
                            .map(AuthUserRole::getRoleId)
                            .distinct()
                            .collect(Collectors.toList());

                    return roleRepository.findAllById(roleIds)
                            .collectList()
                            .map(roles -> {
                                Map<Long, AuthRole> roleMap = roles.stream()
                                        .collect(Collectors.toMap(AuthRole::getId, role -> role));

                                Map<Long, List<AuthRole>> userRoleMap = userRoles.stream()
                                        .collect(Collectors.groupingBy(AuthUserRole::getUserId,
                                                Collectors.mapping(ur -> roleMap.get(ur.getRoleId()),
                                                        Collectors.toList())));

                                return userIds.stream()
                                        .map(userId -> userRoleMap.getOrDefault(userId, List.of()))
                                        .collect(Collectors.toList());
                            });
                });
    }
}
